// STL
#include <iostream>
#include <iomanip>
#include <string>
#include <memory>
#include <cmath>
// ROS
#include <ros/ros.h>
#include <std_msgs/String.h>
#include <geometry_msgs/Twist.h>
// OpenCV
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
// User
#include "common_msg/ObjectPosition.h"
#include "common_msg/TrackedObject.h"
#include "common_utils.hpp"
#include "astra_pro.hpp"
#include "fdsst_tracker.hpp"
#include "kcf_tracker.hpp"
#include "dsst_tracker.hpp"
#include "controller.hpp"

AstraPro astraPro;
cv::Mat originTrackImg(480, 640, CV_8UC3);
cv::Rect2d originTrackRect(-1, -1, -1, -1);

cv::Point dragStartPoint(-1, -1);
cv::Point selectedPoint(-1, -1);
cv::Rect2d selectedRect(-1, -1, -1, -1); // tlx, tly, w, h
cv::Rect2d curTrackRect(-1, -1, -1, -1);
bool objectSelected = false;
cv::Scalar scalarRed(0, 0, 255);
cv::Scalar scalarGreen(0, 255, 0);
cv::Scalar scalarYellow(0, 255, 255);

void setDragStartPoint(int x, int y)
{
	dragStartPoint.x = x;
	dragStartPoint.y = y;
}

void setSelectedPoint(int x, int y)
{
	selectedPoint.x = x;
	selectedPoint.y = y;
}

void setSelectedRect(int x, int y, int w, int h)
{
	selectedRect.x = x;
	selectedRect.y = y;
	selectedRect.width = w;
	selectedRect.height = h;
}

void onMouseClick(int event, int x, int y, int flags, void *userdata)
{
	static int xmin, ymin, xmax, ymax;
	// cout << event << ", "<< x << ", " << y << endl;
	switch (event)
	{
	case cv::EVENT_LBUTTONDOWN:
		if (objectSelected == false)
		{
			setDragStartPoint(x, y);
			setSelectedPoint(x, y);
		}
		break;
	case cv::EVENT_LBUTTONUP:
		setDragStartPoint(-1, -1);
		setSelectedPoint(-1, -1);
		if (objectSelected == false)
		{
			if (selectedRect.width > 60)
			{
				originTrackImg = astraPro.mColorRawImg.clone();
				originTrackRect = selectedRect;
				objectSelected = true;
			}
			else
			{
				setSelectedRect(-1, -1, -1, -1);
			}
		}
		// std::cout << "Depth at " << x << ", " << y << " is " << depth_img.at<ushort>(y / 2, x / 2) << std::endl;
		break;
	case cv::EVENT_MOUSEMOVE:
		if (dragStartPoint.x != -1 && objectSelected == false)
		{
			xmin = std::max(1, std::min(x, dragStartPoint.x));
			ymin = std::max(1, std::min(y, dragStartPoint.y));
			xmax = std::min(639, std::max(x, dragStartPoint.x));
			ymax = std::min(479, std::max(y, dragStartPoint.y));
			setSelectedRect(xmin, ymin, xmax - xmin, ymax - ymin);
		}
		break;
	default:
		break;
	}
}

void drawSelection()
{
	if (dragStartPoint.x != -1 && selectedRect.width != -1)
	{
		cv::rectangle(astraPro.mColorDisImg, selectedRect, scalarRed, 2);
		selectedPoint.x = -1;
	}
	else if (selectedPoint.x != -1)
	{
		cv::circle(astraPro.mColorDisImg, selectedPoint, 10, scalarRed, 2);
	}
}

int selectObject()
{
	const char *title1 = "select object to track";
	const char *title2 = "object to be tracked";
	const char *prompt = "Press any key to continue";
	cv::Point promptPoint(10, 30);

	cv::destroyAllWindows();
	setDragStartPoint(-1, -1);
	setSelectedPoint(-1, -1);
	setSelectedRect(-1, -1, -1, -1);
	objectSelected = false;

	cv::namedWindow(title1, cv::WINDOW_AUTOSIZE);
	cv::setMouseCallback(title1, onMouseClick, NULL);
	int exitFlag = 0;
	while (true)
	{
		astraPro.getColorImg();
		drawSelection();
		cv::imshow(title1, astraPro.mColorDisImg);
		if (cv::waitKey(10) == 27) // press ESC to exit
		{
			exitFlag = -1;
			break;
		}
		if (objectSelected == true)
		{
			cv::destroyWindow(title1);
			astraPro.mColorDisImg = originTrackImg.clone();
			cv::rectangle(astraPro.mColorDisImg, originTrackRect, scalarRed, 2);
			cv::putText(astraPro.mColorDisImg, prompt, promptPoint, cv::FONT_HERSHEY_SIMPLEX, 1, scalarRed, 2);
			cv::imshow(title2, astraPro.mColorDisImg);
			cv::waitKey(1000);
			exitFlag = 1;
			break;
		}
	}
	cv::destroyAllWindows();
	return exitFlag;
}

int main(int argc, char *argv[])
{
	// initialize tracker node
	ros::init(argc, argv, "tracker_node");
	ros::NodeHandle nh("~");
	image_transport::ImageTransport it(nh);
	image_transport::Publisher imgPub = it.advertise("tracked_frame", 5);
	ros::Publisher objPub = nh.advertise<common_msg::TrackedObject>("tracked_object", 5);
	ros::Publisher cmdPub = nh.advertise<geometry_msgs::Twist>("cmd_vel", 5);

	// load param
	std::string frameId, trackerType;
	double thresh, fxy;
	int cameraId, pubRate;
	bool pubCmdVel, pubTrackedFrame, showTrackedFrame;
	nh.param<std::string>("frame_id", frameId, "camera_link");
	nh.param<std::string>("tracker_type", trackerType, "fdsst");
	nh.param<double>("fxy", fxy, 1.0);
	nh.param<int>("camera_id", cameraId, 0);
	nh.param<int>("pub_rate", pubRate, 24);
	nh.param<bool>("pub_cmd_vel", pubCmdVel, true);
	nh.param<bool>("pub_tracked_frame", pubTrackedFrame, false);
	nh.param<bool>("show_tracked_frame", showTrackedFrame, true);
	Controller controller;
	nh.param<int>("track_distance", controller.trackDistance(), 1200);
	nh.param<double>("max_linear_x", controller.maxLinearX(), 0.3);
	nh.param<double>("max_angular_z", controller.maxAngularZ(), 0.3);
	nh.param<double>("linear_p", controller.linearP(), 0.0004);
	nh.param<double>("angular_p", controller.angularP(), 0.0004);

	// initialize camera
	int result = astraPro.open(cameraId);
	if (result == 0)
	{
		astraPro.showInfo();
	}
	else
	{
		std::cout << "Initialize camera failed, return " << result << std::endl;
		return result;
	}

	// tracker data
	sensor_msgs::ImagePtr imgPtr;
	std_msgs::Header header;
	header.frame_id = frameId;
	uint64_t seq = 0;
	common_msg::TrackedObject obj;
	geometry_msgs::Twist trackCmdVel;
	std::shared_ptr<BaseTracker> tracker;
	std::cout << "Tracker type: " << trackerType << std::endl;
	if (trackerType.compare("fdsst") == 0)
	{
		tracker = std::make_shared<FdsstTracker>();
	}
	else if (trackerType.compare("kcf") == 0)
	{
		tracker = std::make_shared<KcfTracker>();
	}
	else if (trackerType.compare("dsst") == 0)
	{
		tracker = std::make_shared<DsstTracker>();
	}
	else
	{
		tracker = std::make_shared<FdsstTracker>();
	}

	// ros loop
	char state = 0;
	char cpsStr[12];
	CpsTimer cpsTimer;
	cv::Point cpsPoint(10, 25);
	ros::Rate loopRate(pubRate);
	int key = 0;
	while (ros::ok())
	{
		switch (state)
		{
		case 0: // show image and select an object to track
		{
			state = selectObject();
			break;
		}
		case 1: // initialize tracker
		{
			tracker->init(originTrackImg, originTrackRect);
			state = 2;
			cpsTimer.start();
			break;
		}
		case 2: // run tracker
		{
			astraPro.getColorImg();
			astraPro.getDepthImg();
			try
			{
				tracker->update(astraPro.mColorRawImg, curTrackRect);
			}
			catch (std::exception &e)
			{
				std::cout << e.what() << std::endl;
				state = 0;
				continue;
			}

			// update header
			header.seq = seq++;
			header.stamp = ros::Time::now();

			// publish tracked object
			obj.header = header;
			obj.is_missing = false;
			astraPro.calcObjPosition(obj, curTrackRect);
			objPub.publish(obj);

			// publish cmd_vel
			if (pubCmdVel)
			{
				controller.calcCmdVel(obj, trackCmdVel);
				cmdPub.publish(trackCmdVel);
			}

			// process tracked frame
			cv::rectangle(astraPro.mColorDisImg, curTrackRect, scalarYellow, 2);
			sprintf(cpsStr, "FPS: %2.1f", cpsTimer.getAvgCps());
			cv::putText(astraPro.mColorDisImg, cpsStr, cpsPoint, cv::FONT_HERSHEY_SIMPLEX, 0.8, scalarYellow, 1);
			if (fxy < 0.9 || fxy > 1.1)
			{
				cv::resize(astraPro.mColorDisImg, astraPro.mColorDisImg, cv::Size(), fxy, fxy);
			}

			// publish tracked frame topic
			if (pubTrackedFrame)
			{
				imgPtr = cv_bridge::CvImage(header, "bgr8", astraPro.mColorDisImg).toImageMsg();
				imgPub.publish(imgPtr);
			}

			// show tracked frame on opencv window
			if (showTrackedFrame)
			{
				cv::imshow("tracked_frame", astraPro.mColorDisImg);
				key = cv::waitKey(10);
				if (key == 'r') // re-select object
				{
					state = 0;
				}
				if (key == 27)
				{
					state = -1;
				}
			}
			break;
		}
		case -1:
		{
			state = -2;
			ros::shutdown();
			break;
		}
		default:
			break;
		}

		ros::spinOnce();
		loopRate.sleep();
	}

	cv::destroyAllWindows();

	astraPro.close();

	return 0;
}